#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
测试Colossal AI Monitor的基本功能。
"""

import os
import json
import shutil
import unittest
import tempfile
import torch
import torch.nn as nn
import torch.optim as optim

from colo_monitor import TrainerMon


class SimpleModel(nn.Module):
    """简单的测试模型."""
    
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(10, 20)
        self.layer2 = nn.Linear(20, 1)
        self.activation = nn.ReLU()
    
    def forward(self, x):
        x = self.activation(self.layer1(x))
        return self.layer2(x)


class TestMonitor(unittest.TestCase):
    """测试监控器的基本功能."""
    
    def setUp(self):
        # 创建临时目录
        self.test_dir = tempfile.mkdtemp()
        os.environ["MONITOR_OUTPUT_DIR"] = self.test_dir
        
        # 创建测试配置文件
        self.config_path = os.path.join(self.test_dir, "monitor_config.json")
        self.create_test_config()
        
        # 创建模型和优化器
        self.model = SimpleModel()
        self.optimizer = optim.Adam(self.model.parameters(), lr=0.001)
        
        # 创建测试数据
        self.inputs = torch.randn(10, 10)
        self.targets = torch.randn(10, 1)
    
    def tearDown(self):
        # 清理临时目录
        shutil.rmtree(self.test_dir)
    
    def create_test_config(self):
        """创建测试配置文件."""
        config = {
            "targets": {},  # 监控所有参数
            "dynamic_on": False,
            "start_step": 0,
            "collect_times": 10,
            "step_interval": 1,
            "print_struct": False,
            "wg_distribution": True,
            "param_distribution": True,
            "xy_distribution": False,
            "mv_distribution": False,
            "format": "csv",
            "ops": ["norm", "min", "max", "mean"],
            "eps": 1e-8,
            "ndigits": 6,
            "step_count_per_record": 1,
            "squash_name": True
        }
        
        with open(self.config_path, "w") as f:
            json.dump(config, f, indent=4)
    
    def test_basic_monitoring(self):
        """测试基本监控功能."""
        # 创建监控器
        monitor = TrainerMon(config_file_path=self.config_path)
        
        # 挂载监控对象
        monitor.set_monitor(
            self.model,
            grad_acc_steps=1,
            optimizer=self.optimizer
        )
        
        # 简单训练
        criterion = nn.MSELoss()
        
        for i in range(5):
            # 前向传播
            outputs = self.model(self.inputs)
            loss = criterion(outputs, self.targets)
            
            # 反向传播
            self.optimizer.zero_grad()
            loss.backward()
            monitor.process_after_backward()
            
            # 更新参数
            self.optimizer.step()
            monitor.process_after_step()
        
        # 获取和检查梯度指标
        reduced_metrics, unreduced_metrics = monitor.generate_wgrad_metrics()
        
        # 确保有指标被收集
        self.assertGreater(len(reduced_metrics), 0)
        self.assertGreater(len(unreduced_metrics), 0)
        
        # 检查指标内容
        for name, metrics in reduced_metrics.items():
            self.assertIn("norm", metrics)
            self.assertIn("min", metrics)
            self.assertIn("max", metrics)
            self.assertIn("mean", metrics)
        
        # 关闭监控器
        monitor.close()
        
        # 检查输出文件是否生成
        files = os.listdir(self.test_dir)
        csv_files = [f for f in files if f.endswith('.csv')]
        self.assertGreater(len(csv_files), 0)
    
    def test_selective_monitoring(self):
        """测试选择性监控功能."""
        # 创建选择性监控配置
        config = {
            "targets": {"layer1": {}},  # 只监控layer1
            "wg_distribution": True,
            "format": "csv",
            "ops": ["norm"]
        }
        
        selective_config_path = os.path.join(self.test_dir, "selective_config.json")
        with open(selective_config_path, "w") as f:
            json.dump(config, f, indent=4)
        
        # 创建监控器
        monitor = TrainerMon(config_file_path=selective_config_path)
        
        # 挂载监控对象
        monitor.set_monitor(
            self.model,
            grad_acc_steps=1,
            optimizer=self.optimizer
        )
        
        # 简单训练
        criterion = nn.MSELoss()
        outputs = self.model(self.inputs)
        loss = criterion(outputs, self.targets)
        self.optimizer.zero_grad()
        loss.backward()
        monitor.process_after_backward()
        self.optimizer.step()
        monitor.process_after_step()
        
        # 获取和检查梯度指标
        reduced_metrics, _ = monitor.generate_wgrad_metrics()
        
        # 确保只有layer1的参数被监控
        for name in reduced_metrics.keys():
            self.assertIn("layer1", name)
        
        # 关闭监控器
        monitor.close()
    
    def test_api_format(self):
        """测试API格式输出."""
        # 创建API格式监控配置
        config = {
            "targets": {},
            "wg_distribution": True,
            "format": "api",  # 使用API格式
            "ops": ["norm"]
        }
        
        api_config_path = os.path.join(self.test_dir, "api_config.json")
        with open(api_config_path, "w") as f:
            json.dump(config, f, indent=4)
        
        # 创建监控器
        monitor = TrainerMon(config_file_path=api_config_path)
        
        # 挂载监控对象
        monitor.set_monitor(
            self.model,
            grad_acc_steps=1,
            optimizer=self.optimizer
        )
        
        # 简单训练
        criterion = nn.MSELoss()
        outputs = self.model(self.inputs)
        loss = criterion(outputs, self.targets)
        self.optimizer.zero_grad()
        loss.backward()
        monitor.process_after_backward()
        self.optimizer.step()
        monitor.process_after_step()
        
        # 获取和检查梯度指标
        reduced_metrics, _ = monitor.generate_wgrad_metrics()
        
        # 确保有指标被收集
        self.assertGreater(len(reduced_metrics), 0)
        
        # 检查指标内容
        for metrics in reduced_metrics.values():
            self.assertIn("norm", metrics)
        
        # 关闭监控器
        monitor.close()
        
        # 检查API格式不生成CSV文件
        files = os.listdir(self.test_dir)
        csv_files = [f for f in files if f.endswith('.csv') and "api_config" not in f]
        self.assertEqual(len(csv_files), 0)


if __name__ == "__main__":
    unittest.main()
